import time
import numpy as np
import pandas as pd
import warnings
warnings.filterwarnings("ignore")


def backtest(all_buys):
    stocks_path = r'd:\Python\study_data\stocks\\'
    all_buys = pd.read_csv(all_buys)
    s_5 = []
    s_10 = []
    s_20 = []
    s_25 = []
    l_5 = []
    l_10 = []
    l_20 = []
    l_25 = []
    cnt = 0
    win_5 = 0
    win_10 = 0
    win_20 = 0
    win_25 = 0
    for i in range(len(all_buys)):
        row = list(all_buys.iloc[i].dropna())
        date = row[0]
        codes = row[1:]
        for code in codes:
            stock_df = pd.read_csv(stocks_path + code + '.csv', index_col=0)
            day = stock_df.index.get_loc(date)
            buy = stock_df.iloc[day]['close']
            if day + 25 < len(stock_df):
                sell_5 = round(stock_df.iloc[day + 5]['close'] / buy - 1, 2)
                sell_10 = round(stock_df.iloc[day + 10]['close'] / buy - 1, 2)
                sell_20 = round(stock_df.iloc[day + 20]['close'] / buy - 1, 2)
                sell_25 = round(stock_df.iloc[day + 25]['close'] / buy - 1, 2)
                if np.isinf(sell_5) or np.isinf(sell_10) or np.isinf(sell_20) or np.isinf(sell_25) or np.isnan(sell_5) or np.isnan(sell_10) or np.isnan(sell_20) or np.isnan(sell_25):
                    continue
                else:
                    if sell_5 > 0:
                        win_5 += 1
                        s_5.append(sell_5)
                    else:
                        l_5.append(sell_5)
                    if sell_10 > 0:
                        win_10 += 1
                        s_10.append(sell_10)
                    else:
                        l_10.append(sell_10)
                    if sell_20 > 0:
                        win_20 += 1
                        s_20.append(sell_20)
                    else:
                        l_20.append(sell_20)
                    if sell_25 > 0:
                        win_25 += 1
                        s_25.append(sell_25)
                    else:
                        l_25.append(sell_25)
                    cnt += 1
            # break
        # break

    pnt = round(win_5 / cnt, 2)
    print(pnt, np.mean(s_5), np.mean(l_5), pnt * np.mean(s_5) + (1 - pnt) * np.mean(l_5))
    pnt = round(win_10 / cnt, 2)
    print(pnt, np.mean(s_10), np.mean(l_10), pnt * np.mean(s_10) + (1 - pnt) * np.mean(l_10))
    pnt = round(win_20 / cnt, 2)
    print(pnt, np.mean(s_20), np.mean(l_20), pnt * np.mean(s_20) + (1 - pnt) * np.mean(l_20))
    pnt = round(win_25 / cnt, 2)
    print(pnt, np.mean(s_25), np.mean(l_25), pnt * np.mean(s_25) + (1 - pnt) * np.mean(l_25))

    return None


if __name__ == '__main__':
    begin_time = time.time()
    test = 'all_buys_20_day_line.csv'
    backtest(test)

    end_time = time.time()
    run_time = round(end_time - begin_time)
    hour = run_time // 3600
    minute = (run_time - 3600 * hour) // 60
    second = run_time - 3600 * hour - 60 * minute
    print(f'该程序运行时间：{hour}小时{minute}分钟{second}秒')


